﻿#pragma once
#include "DataObject.h"
#include <vtkm/cont/DataSetBuilderUniform.h>
#include <vtkm/worklet/WorkletMapTopology.h>

class MeshBase
{
public:
  struct ComputeTopolgy : vtkm::worklet::WorkletVisitCellsWithPoints
  {
    using ControlSignature = void(CellSetIn, FieldInPoint, FieldOut, FieldOut);
    using ExecutionSignature = void(_2, _3, _4);

    template<typename PType>
    VTKM_EXEC void operator()(const PType& pts, vtkm::Vec3f& centroid, vtkm::Vec3f& spacing) const
    {
      if (pts.GetNumberOfComponents() == 4)
      {
        centroid = (pts[0] + pts[1] + pts[2] + pts[3]) / 4;
        spacing[0] = (pts[1] - pts[0])[0];
        spacing[1] = (pts[2] - pts[1])[1];
      }
      else if (pts.GetNumberOfComponents() == 8)
      {
        centroid = (pts[0] + pts[1] + pts[2] + pts[3] + pts[4] + pts[5] + pts[6] + pts[7]) / 8;
        spacing[0] = (pts[1] - pts[0])[0];
        spacing[1] = (pts[2] - pts[1])[1];
        spacing[2] = (pts[4] - pts[0])[2];
      }
      else
      {
        RaiseError("unfinished");
      }
    }
  };

public:
  MeshBase()
  {
  }
  MeshBase(vtkm::Id2 dim, vtkm::Vec2f start, vtkm::Vec2f end)
  {
    this->Create(dim, start, end);
    vtkm::cont::Invoker{}(ComputeTopolgy{}, _cell_set, _coord, _centroid, _spacing);
  }
  MeshBase(vtkm::Id3 dim, vtkm::Vec3f start, vtkm::Vec3f end)
  {
    this->Create(dim, start, end);
    vtkm::cont::Invoker{}(ComputeTopolgy{}, _cell_set, _coord, _centroid, _spacing);
  }
  virtual ~MeshBase(){};

  void Create(vtkm::Id3 dim, vtkm::Vec3f start, vtkm::Vec3f end)
  {
    auto space = end - start;
    vtkm::Id3 point_dim = { dim[0] + 1, dim[1] + 1, dim[2] + 1 };

    if (dim[1] > 0)
    {
      for (size_t i = 0; i < 2; i++)
        space[i] = (end[i] - start[i]) / dim[i];

      space[2] = 1;
      vtkm::cont::DataSet result =
        vtkm::cont::DataSetBuilderUniform::Create(point_dim, start, space);

      this->_cell_set = result.GetCellSet();
      this->_coord = result.GetCoordinateSystem();
    }
  }

  void Create(vtkm::Id2 dim, vtkm::Vec2f start, vtkm::Vec2f end)
  {
    vtkm::Vec2f space = end - start;
    vtkm::Id2 point_dim = { dim[0] + 1, dim[1] + 1 };

    for (size_t i = 0; i < 2; i++)
      space[i] = (end[i] - start[i]) / dim[i];


    vtkm::cont::DataSet result = vtkm::cont::DataSetBuilderUniform::Create(point_dim, start, space);

    this->_cell_set = result.GetCellSet();
    this->_coord = result.GetCoordinateSystem();
  }

  CoordinateSystem& GetCoordinate() { return this->_coord; }
  const CoordinateSystem& GetCoordinate() const { return this->_coord; }

  const vtkm::cont::DynamicCellSet& GetCellSet() const { return this->_cell_set; }
  vtkm::cont::DynamicCellSet& GetCellSet() { return this->_cell_set; }

private:
  DynamicCellSet _cell_set;
  CoordinateSystem _coord;

public:
  ArrayHandle<vtkm::Vec3f> _centroid;
  ArrayHandle<vtkm::Vec3f> _spacing;
};
